import os
import requests
import json
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Get API key from environment variables
API_KEY = os.getenv("DEEPSEEK_API_KEY")
if not API_KEY or API_KEY == "your_api_key_here":
    print("Please set your DeepSeek API key in the .env file")
    exit(1)

# DeepSeek API endpoint
API_URL = "https://api.deepseek.com/v1/chat/completions"

def stream_response(prompt, model="deepseek-chat", max_tokens=1000):
    """
    Stream a response from the DeepSeek API
    
    Args:
        prompt (str): The user's input prompt
        model (str): The model to use (default: deepseek-chat)
        max_tokens (int): Maximum number of tokens to generate
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    data = {
        "model": model,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "max_tokens": max_tokens,
        "stream": True  # Enable streaming
    }
    
    try:
        response = requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True)
        response.raise_for_status()
        
        # Process the streaming response
        for line in response.iter_lines():
            if line:
                # Remove the "data: " prefix and parse JSON
                line_text = line.decode('utf-8')
                if line_text.startswith('data: '):
                    json_str = line_text[6:]  # Remove 'data: ' prefix
                    
                    # Skip [DONE] message
                    if json_str.strip() == "[DONE]":
                        break
                    
                    try:
                        chunk = json.loads(json_str)
                        if "choices" in chunk and len(chunk["choices"]) > 0:
                            delta = chunk["choices"][0].get("delta", {})
                            if "content" in delta:
                                print(delta["content"], end="", flush=True)
                    except json.JSONDecodeError:
                        pass
        
        print()  # Add a newline at the end
    except requests.exceptions.RequestException as e:
        print(f"\nError calling DeepSeek API: {e}")

def main():
    print("DeepSeek API Streaming Chat Interface")
    print("Type 'exit' to quit")
    print("-" * 50)
    
    while True:
        user_input = input("\nYou: ")
        
        if user_input.lower() in ["exit", "quit"]:
            print("Goodbye!")
            break
        
        print("\nDeepSeek: ", end="", flush=True)
        stream_response(user_input)

if __name__ == "__main__":
    main() 